import numpy as np
import cv2
import pickle
import matplotlib.pyplot as plt
from create_clf import create_features

def main():
    # 加载模型
    with open('sandstone_clf.pkl', 'rb') as f:
        clf = pickle.load(f)
    
    # 读取图像
    image = cv2.imread('Sandstone_imgs/Sandstone_2.tif', cv2.IMREAD_GRAYSCALE)
    true_segment = cv2.imread('Sandstone_imgs/Sandstone_2_segment.tif', cv2.IMREAD_GRAYSCALE)
    
    # 创建特征
    X = create_features(image)
    
    # 预测
    predictions = clf.predict(X)
    predicted_image = predictions.reshape(image.shape)
    
    # 计算准确率
    accuracy = np.mean(predicted_image.ravel() == true_segment.ravel())
    
    # 显示结果
    plt.figure(figsize=(15, 5))
    
    plt.subplot(131)
    plt.imshow(image, cmap='gray')
    plt.title('original image')
    plt.axis('off')
    
    plt.subplot(132)
    plt.imshow(true_segment, cmap='gray')
    plt.title('segment')
    plt.axis('off')
    
    plt.subplot(133)
    plt.imshow(predicted_image, cmap='gray')
    plt.title(f'segmentation(acc:{accuracy:.3f})')
    plt.axis('off')
    
    plt.tight_layout()
    plt.show()

if __name__ == "__main__":
    main() 